import json

import uuid

import boto3

from framework.common.pt_logger import logger
from framework.core.crawler_result import Result


class SQSOutput(object):
    def __init__(self, queue_name, get_queue_message=None, condition=None):
        self.queue_name = queue_name
        self.get_queue_message = get_queue_message
        self.condition = condition
        self.client = boto3.client('sqs')

    @staticmethod
    def done():
        logger.info("output done")
        pass

    def _dump_single(self, current_task, result):
        if (self.condition is None) or (self.condition is not None and self.condition(current_task, result)):
            queue_messages = self.get_queue_message(current_task, result)
            logger.info("output to sqs, queue_messages are %s" % queue_messages)
            if isinstance(queue_messages, list):
                uuids = [str(uuid.uuid1()) for _ in queue_messages]
                entries = list()
                for i, t in enumerate(queue_messages):
                    task_obj = {
                        "Id": uuids[i],
                        "MessageBody": queue_messages[i]
                    }
                    entries.append(task_obj)
                # JimmyMo: 一次性可以入队10个
                entries_batches = [entries[x:x + 10] for x in range(0, len(entries), 10)]
                for entries_batch in entries_batches:
                    response = boto3.client('sqs').send_message_batch(
                        QueueUrl="https://sqs.us-east-2.amazonaws.com/599885534814/%s" % self.queue_name,
                        Entries=entries_batch
                    )
                    logger.info("sqs response: %s" % response)
                    if response.get("ResponseMetadata", dict()).get("HTTPStatusCode", None) == 200:
                        logger.info('output file to %s...' % self.queue_name)
                    else:
                        logger.warning("failed to output file to %s..." % self.queue_name)
            else:
                queue_message = queue_messages
                response = self.client.send_message(
                    QueueUrl="https://sqs.us-east-2.amazonaws.com/599885534814/%s" % self.queue_name,
                    MessageBody=queue_message,
                )
                logger.info("sqs response: %s" % response)
                if response.get("ResponseMetadata", dict()).get("HTTPStatusCode", None) == 200:
                    logger.info('output file to %s...' % self.queue_name)
                else:
                    logger.warning("failed to output file to %s..." % self.queue_name)

    def dump(self, current_task, results):
        if isinstance(results, Result):
            result = results
            self._dump_single(current_task, result)
        elif isinstance(results, list) and len(results) > 0:
            for result in results:
                self._dump_single(current_task, result)


if __name__ == "__main__":
    pass
